// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/unit.h>
#include <melon/executors/queued_immediate_executor.h>
#include <melon/futures/future.h>
#include <functional>
#include <memory>
#include <pollux/common/time/cpu_wall_timer.h>

#include <pollux/common/base/exceptions.h>
#include <pollux/common/base/portability.h>
#include <pollux/common/future/pollux_promise.h>
#include <pollux/common/process/thread_debug_info.h>
#include <pollux/common/testutil/test_value.h>

namespace kumo::pollux {
    // A future-like object that prefabricates Items on an executor and
    // allows consumer threads to pick items as they are ready. If the
    // consumer needs the item before the executor started making it, the
    // consumer will make it instead. If multiple consumers request the
    // same item, exactly one gets it. Propagates exceptions to the
    // consumer.
    template<typename Item>
    class AsyncSource {
    public:
        explicit AsyncSource(std::function<std::unique_ptr<Item>()> make)
            : make_(std::move(make)) {
            if (process::GetThreadDebugInfo() != nullptr) {
                auto *currentThreadDebugInfo = process::GetThreadDebugInfo();
                // We explicitly leave out the callback when copying the ThreadDebugInfo
                // as that may have captured state that goes out of scope by the time
                // _make is called.
                threadDebugInfo_ = std::make_optional<process::ThreadDebugInfo>(
                    {
                        currentThreadDebugInfo->queryId_,
                        currentThreadDebugInfo->taskId_,
                        nullptr
                    });
            }
        }

        ~AsyncSource() {
            POLLUX_CHECK(
                moved_ || closed_,
                "AsyncSource should be properly consumed or closed.");
        }

        // Makes an item if it is not already made. To be called on a background
        // executor.
        void prepare() {
            common::testutil::TestValue::adjust(
                "kumo::pollux::AsyncSource::prepare", this);
            std::function<std::unique_ptr<Item>()> make = nullptr; {
                std::lock_guard<std::mutex> l(mutex_);
                if (!make_) {
                    return;
                }
                making_ = true;
                std::swap(make, make_);
            }
            std::unique_ptr<Item> item;
            try {
                CpuWallTimer timer(timing_);
                item = runMake(make);
            } catch (std::exception &) {
                std::lock_guard<std::mutex> l(mutex_);
                exception_ = std::current_exception();
            }
            std::unique_ptr<ContinuePromise> promise; {
                std::lock_guard<std::mutex> l(mutex_);
                POLLUX_CHECK_NULL(item_);
                if (MELON_LIKELY(exception_ == nullptr)) {
                    item_ = std::move(item);
                }
                making_ = false;
                promise.swap(promise_);
            }
            if (promise != nullptr) {
                promise->setValue();
            }
        }

        // Returns the item to the first caller and nullptr to subsequent callers.
        // If the item is preparing on the executor, waits for the item and
        // otherwise makes it on the caller thread.
        std::unique_ptr<Item> move() {
            common::testutil::TestValue::adjust(
                "kumo::pollux::AsyncSource::move", this);
            std::function<std::unique_ptr<Item>()> make = nullptr;
            ContinueFuture wait; {
                std::lock_guard<std::mutex> l(mutex_);
                moved_ = true;
                // 'making_' can be read atomically, 'exception' maybe not. So test
                // 'making' so as not to read half-assigned 'exception_'.
                if (!making_ && exception_) {
                    std::rethrow_exception(exception_);
                }
                if (item_) {
                    return std::move(item_);
                }
                if (promise_) {
                    // Somebody else is now waiting for the item to be made.
                    return nullptr;
                }
                if (making_) {
                    promise_ = std::make_unique<ContinuePromise>();
                    wait = promise_->getSemiFuture();
                } else {
                    if (!make_) {
                        return nullptr;
                    }
                    std::swap(make, make_);
                }
            }
            // Outside of mutex_.
            if (make) {
                try {
                    return runMake(make);
                } catch (const std::exception &) {
                    std::lock_guard<std::mutex> l(mutex_);
                    exception_ = std::current_exception();
                    throw;
                }
            }
            auto &exec = melon::QueuedImmediateExecutor::instance();
            std::move(wait).via(&exec).wait();
            std::lock_guard<std::mutex> l(mutex_);
            if (exception_) {
                std::rethrow_exception(exception_);
            }
            return std::move(item_);
        }

        // If true, move() will not block. But there is no guarantee that somebody
        // else will not get the item first.
        bool hasValue() const {
            tsan_lock_guard<std::mutex> l(mutex_);
            return item_ != nullptr || exception_ != nullptr;
        }

        /// Returns the timing of prepare(). If the item was made on the calling
        /// thread, the timing is 0 since only off-thread activity needs to be added
        /// to the caller's timing.
        const CpuWallTiming &prepareTiming() {
            return timing_;
        }

        /// This function assists the caller in ensuring that resources allocated in
        /// AsyncSource are promptly released:
        /// 1. Waits for the completion of the 'make_' function if it is executing
        /// in the thread pool.
        /// 2. Resets the 'make_' function if it has not started yet.
        /// 3. Cleans up the 'item_' if 'make_' has completed, but the result
        /// 'item_' has not been returned to the caller.
        void close() {
            if (closed_ || moved_) {
                return;
            }
            ContinueFuture wait; {
                std::lock_guard<std::mutex> l(mutex_);
                if (making_) {
                    promise_ = std::make_unique<ContinuePromise>();
                    wait = promise_->getSemiFuture();
                } else if (make_) {
                    make_ = nullptr;
                }
            }

            auto &exec = melon::QueuedImmediateExecutor::instance();
            std::move(wait).via(&exec).wait(); {
                std::lock_guard<std::mutex> l(mutex_);
                if (item_) {
                    item_ = nullptr;
                }
                closed_ = true;
            }
        }

    private:
        std::unique_ptr<Item> runMake(std::function<std::unique_ptr<Item>()> &make) {
            process::ScopedThreadDebugInfo threadDebugInfo(
                threadDebugInfo_.has_value() ? &threadDebugInfo_.value() : nullptr);
            return make();
        }

        // Stored context (if present upon construction) so they can be restored when
        // make_ is invoked.
        std::optional<process::ThreadDebugInfo> threadDebugInfo_;

        mutable std::mutex mutex_;
        // True if 'prepare() is making the item.
        bool making_{false};
        std::unique_ptr<ContinuePromise> promise_;
        std::unique_ptr<Item> item_;
        std::function<std::unique_ptr<Item>()> make_;
        std::exception_ptr exception_;
        CpuWallTiming timing_;
        bool closed_{false};
        bool moved_{false};
    };
} // namespace kumo::pollux
